import re
import os as real_os
from server import config
from pydantic import BaseModel
from typing import Literal
from neutron import Plugin, execute, User, OS, Shutil
from fastapi import WebSocket, HTTPException
from fastapi.responses import FileResponse
from starlette import websockets
from datetime import datetime
# import os
import asyncio
# import shutil
import pickle
import stat
import aiofiles
import anyio
import pwd
import grp
import stat
from starlette import types as starlette_types
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from . import icon
from . import filemover
from . import search

app = Plugin(__file__)


async def timeout_second(tasks):  # 防超时
    try:
        return await asyncio.wait_for(asyncio.shield(asyncio.create_task(asyncio.to_thread(tasks))), timeout=1)
    except TimeoutError:
        raise TimeoutError("任务超时")


def format_size(size_bytes: int) -> str:
    if (size_bytes > 1024*1024*1024*1024):
        return f"{size_bytes/(1024*1024*1024*1024):.2f}TiB"
    elif (size_bytes > 1024*1024*1024):
        return f"{size_bytes/(1024*1024*1024):.2f}GiB"
    elif (size_bytes > 1024*1024):
        return f"{size_bytes/(1024*1024):.2f}MiB"
    elif (size_bytes > 1024):
        return f"{size_bytes/(1024):.2f}KiB"
    else:
        return f"{size_bytes}B"


async def get_size(os, path):  # 获取文件大小并添加单位
    if (await os.path.isdir(path)):
        return ""
    try:
        size_bytes = (await os.stat(path)).st_size
        return format_size(size_bytes)
    except:
        return ""


@app.route.get("/files")
async def get_files_list(user: User, path: str, show_hides: int = 0):  # 获取文件列表
    os = OS(user)
    try:
        files = await os.listdir(path)
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}
    except:
        return {"code": 1, "msg": "找不到路径", "data": []}
    try:
        datas_lst = [
            {
                "img": "icons/"+(icon.get_icon(path+os.sep+i)),
                "name": i,
                "size": await get_size(os, path+os.sep+i),
                "isdir": await os.path.isdir(path+os.sep+i)
            } for i in files
        ]
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}

    async def sort_key(x):  # 排序
        # 文件夹优先
        return ("0" if await os.path.isdir(path+os.sep+x["name"]) else "1")+x["name"]
    keys = await (asyncio.gather(
        *(sort_key(item) for item in datas_lst)))

    datas_lst = [x for _, x in (await timeout_second(lambda: sorted(zip(keys, datas_lst))))]
    if (show_hides == 0):  # 是否显示隐藏文件
        datas_lst = [i for i in datas_lst if not i["name"].startswith(".")]
    return {"code": 0, "msg": "", "data": datas_lst}

openers = {
}

bind_exts = {
}


@app.route.get("/open")  # 打开文件
async def open_file(user: User, path: str):
    os = OS(user)
    if (not await os.path.exists(path)):
        return {"code": 1, "msg": path+" 找不到路径", "data": []}
    if (await os.path.isdir(path)):
        return {"code": 1, "msg": path+" 不是一个文件", "data": []}
    ext = path.split(os.sep)[-1].split(".")[-1]
    if (ext in bind_exts):
        return {"code": 0, "msg": "", "type": "open", "data":  openers[bind_exts[ext]]["code"].format(path)}
    return {"code": 0, "msg": "", "type": "choose", "data": {
        k: {
            "code": v["code"].format(path),
            "icon": v["icon"],
            "name": v["name"]
        } for k, v in openers.items()
    }}


@app.export  # 注册文件打开方式
def bind_opener(opener_id: str, opener_code: str, opener_icon: str, opener_name: str):
    openers[opener_id] = {
        "code": opener_code,
        "icon": opener_icon,
        "name": opener_name
    }


@app.export  # 注册扩展名
def bind_ext(ext: str, opener_id: str):
    bind_exts[ext] = opener_id


@app.route.post("/new/file")  # 新建文件
async def new_file(user: User, path: str):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (await os.path.exists(path)):
            return {"code": 1, "msg": "文件已存在", "data": []}
        if (not await os.path.exists(os.path.dirname(path))):
            return {"code": 1, "msg": "目录不存在", "data": []}
        if ('&' in path or '|' in path or '\\' in path or ';' in path or '*' in path or '?' in path or '#' in path):
            return {"code": 1, "msg": "文件名包含非法字符", "data": []}

        async def new_file():
            f: OS.modetext
            async with os.open(path, "w") as f:  # type: ignore
                await f.write("")
        await (new_file())
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}
    except Exception as e:
        return {"code": 1, "msg": "创建失败", "data": []}
    return {"code": 0, "msg": "", "data": []}


@app.route.post("/new/dir")
async def new_dir(user: User, path: str):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (await os.path.exists(path)):
            return {"code": 1, "msg": "文件已存在", "data": []}
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}
    if ('&' in path or '|' in path or '\\' in path or ';' in path or '*' in path or '?' in path or '#' in path or '\'' in path or '"' in path):
        return {"code": 1, "msg": "文件名包含非法字符", "data": []}
    try:
        await os.mkdir(path)
    except:
        return {"code": 1, "msg": "创建失败", "data": []}
    return {"code": 0, "msg": "", "data": []}


@app.route.post("/rename")  # 重命名
async def rename_file_or_dir(user: User, path: str, from_name: str, to_name: str):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (not await os.path.exists(path)):
            return {"code": 1, "msg": "文件不存在", "data": []}
        if (not await os.path.exists(path+os.sep+from_name)):
            return {"code": 1, "msg": "文件不存在", "data": []}
        if (await os.path.exists(path+os.sep+to_name)):
            return {"code": 1, "msg": "文件已存在", "data": []}
        if ('&' in from_name or '|' in from_name or '\\' in from_name or ';' in from_name or '*' in from_name or '?' in from_name or '#' in from_name or '\'' in from_name or '"' in from_name):
            return {"code": 1, "msg": "文件名包含非法字符", "data": []}
        if ('&' in to_name or '|' in to_name or '\\' in to_name or ';' in to_name or '*' in to_name or '?' in to_name or '#' in to_name or '\'' in to_name or '"' in to_name):
            return {"code": 1, "msg": "文件名包含非法字符", "data": []}
        await os.rename(path+os.sep+from_name, path+os.sep+to_name)
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}
    except:
        return {"code": 1, "msg": "重命名失败", "data": []}
    return {"code": 0, "msg": "", "data": []}


class MoveFileModel(BaseModel):
    types: Literal["move", "copy"]
    origin: list[str]
    dist: str


@app.route.post("/paste")  # 粘贴
async def paste_file(user: User, move_file: MoveFileModel):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    if (len(move_file.origin) == 0):
        return {"code": 1, "msg": "没有选择文件", "data": []}
    if (len(move_file.origin) == 1):
        files_info = move_file.origin[0]
    else:
        files_info = f"{len(move_file.origin)} 个文件"
    try:
        if (not await os.path.exists(move_file.dist)):
            return {"code": 1, "msg": "目标文件夹不存在", "data": []}

        if (not await os.access(move_file.dist, os.W_OK)):
            return {
                "code": 1,
                "msg": "权限不足",
                "data": []
            }
    except PermissionError:
        return {
            "code": 1,
            "msg": "权限不足",
            "data": []
        }

    async def move_task(tasktool):
        nonlocal move_file
        await asyncio.sleep(0.1)

        async def write_log(log):
            await (tasktool.write_log(log))

        async def write_progress(progress):
            await (tasktool.set_process(progress))

        async def exit_process(err_flag):
            await (tasktool.exit(err_flag))

        async def wait_button(btn_list):
            return await (tasktool.show_buttons(btn_list))
        await filemover.move_function(move_file, write_log,
                                      write_progress, exit_process, wait_button)
    await execute("tasks", "create_task")(
        ("移动" if move_file.types == "move" else "复制"), files_info+" 到 "+move_file.dist, move_task, user)
    return {"code": 0, "msg": "", "data": []}


class RemoveFileModel(BaseModel):
    origin: list[str]
    types: Literal["delete"]


@app.route.post("/remove")  # 删除
async def remove_file(user: User, move_file: RemoveFileModel):
    os = OS(user)
    shutil = Shutil(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    if (len(move_file.origin) == 0):
        return {"code": 1, "msg": "没有选择文件", "data": []}

    async def remove_task(tasktool):
        nonlocal move_file
        await asyncio.sleep(0.1)

        await tasktool.write_log("删除以下文件：\n")
        for i in move_file.origin:
            await tasktool.write_log(i+"\n")
        await tasktool.write_log("\n是否确认删除，删除后将无法还原！\n")

        res = await tasktool.show_buttons(["确定", "取消"])
        if (res != 0):
            await tasktool.exit(True)
            return

        now_error_str = ""
        err_mode = False

        async def rmtree(dist):
            nonlocal now_error_str
            try:
                if (await os.path.isdir(dist)):
                    await shutil.rmtree(dist)
                else:
                    await os.remove(dist)
            except PermissionError:
                now_error_str = "权限不足："+dist
            except Exception as e:
                now_error_str = "其它错误["+repr(e)+"]："+dist

        j = 0
        for i in move_file.origin:
            j += 1
            await tasktool.set_process(j*100//len(move_file.origin))
            await rmtree(i)
            if (now_error_str != ""):
                await tasktool.write_log(now_error_str+"\n")
                err_mode = True

        await tasktool.set_process(100)
        await tasktool.exit(err_mode)
    await execute("tasks", "create_task")(
        ("删除"), (move_file.origin[0] if len(move_file.origin) == 1 else str(len(move_file.origin)) + " 个文件"), remove_task, user)
    return {"code": 0, "msg": "", "data": []}


@app.route.ws("/search", level=1)
async def search_file(user: User, websocket: WebSocket, search_path: str):
    os = OS(user)
    if (await user.check_permission(search_path, 1, True) == False):
        raise HTTPException(401)
    return await search.search_ws(user, os, app, websocket, search_path)


class BookmarkReceive(BaseModel):
    data: list[dict]
    type: Literal["sync"]


@app.route.post("/bookmark", level=1)
async def save_bookmark(user: User, bookmarks: BookmarkReceive):
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    bookmark_save = await app.filestorage.load(user, "bookmarks")
    async with bookmark_save.open("bookmark.pickle", "wb") as file:
        await file.write(await asyncio.to_thread(lambda: pickle.dumps(bookmarks.data)))
    return {"code": 0, "msg": ""}


@app.route.get("/bookmark")
async def get_bookmark(user: User):
    if (user.demo):
        return {"code": 0, "msg": "", "data": [
            {
                "text": "[DEMO MODE]",
                "path": "/"
            },
            {
                "text": "[root]",
                "path": "/"
            },
            {
                "text": "app",
                "path": "/app"
            }
        ], "new": True}
    bookmark_save = await app.filestorage.load(user, "bookmarks")
    if (not await bookmark_save.exists("bookmark.pickle")):
        return {"code": 0, "msg": "", "data": [{
            "text": "[root]",
            "path": "/"
        }], "new": True}
    async with bookmark_save.open("bookmark.pickle", "rb") as file:
        load_data = await file.read()
        bookmark_lst = await asyncio.to_thread(lambda: pickle.loads(load_data))
    return {"code": 0, "msg": "", "data": bookmark_lst, "new": False}


class WatchPathHandler(FileSystemEventHandler):

    def __init__(self, sendmsg):
        self.send_msg = sendmsg

    def on_created(self, event):
        self.send_msg()

    def on_deleted(self, event):
        self.send_msg()

    def on_moved(self, event):
        self.send_msg()


observer = Observer()
observer.start()


@app.route.ws("/watch", level=1)
async def watch_path(user: User, websocket: WebSocket, path: str):
    os = OS(user)
    try:
        if (not await os.path.exists(path)):
            return {"code": 1, "msg": "路径不存在", "data": []}
    except PermissionError:
        return {"code": 1, "msg": "权限不足", "data": []}
    except:
        return {"code": 1, "msg": "未知错误", "data": []}
    global observer

    event_handler = None
    schedule = None
    tasks = None

    try:
        await websocket.accept()
        update_msg = 0

        def change_callback():
            nonlocal update_msg
            update_msg = 1

        event_handler = WatchPathHandler(change_callback)
        schedule = observer.schedule(event_handler, path, recursive=False)

        async def wait_task():
            nonlocal update_msg
            while True:
                if (update_msg == 1):
                    await websocket.send_json({"code": 0, "msg": "", "data": "change"})
                    update_msg = 0
                await asyncio.sleep(1)

        tasks = asyncio.create_task(wait_task())
        while True:
            await websocket.receive_json()

    except (RuntimeError, websockets.WebSocketDisconnect):
        if (event_handler is not None):
            del event_handler
        if (schedule is not None):
            try:
                observer.unschedule(schedule)
            except KeyError:
                pass
            await asyncio.sleep(0.1)
        if (tasks is not None and tasks.done() == False):
            tasks.cancel()


class RangeFileResponse(FileResponse):
    async def __call__(self, scope: starlette_types.Scope, receive: starlette_types.Receive, send: starlette_types.Send) -> None:
        if self.stat_result is None:
            try:
                stat_result = await timeout_second(lambda: real_os.stat(self.path))
                self.set_stat_headers(stat_result)
            except FileNotFoundError:
                raise RuntimeError(f"File at path {self.path} does not exist.")
            else:
                mode = stat_result.st_mode
                if not stat.S_ISREG(mode):
                    raise RuntimeError(
                        f"File at path {self.path} is not a file.")

        scope_h = {m[0].decode('utf-8'): m[1].decode('utf-8')
                   for m in scope['headers']}
        range_header = scope_h.get('range')
        start, stop = 0, None
        if range_header:
            range_str = range_header
            if range_str.startswith('bytes='):
                range_parts = range_str[6:].split('-')
                if len(range_parts) == 2:
                    start = int(range_parts[0])
                    if range_parts[1]:
                        stop = int(range_parts[1]) + 1
                elif len(range_parts) == 1:
                    start = int(range_parts[0])
            if start < 0 or (stop is not None and stop <= start):
                raise RuntimeError("Invalid Range header")
            self.status_code = 206
            if stop:
                if (stop > stat_result.st_size):
                    stop = stat_result.st_size
                self.headers['Content-Range'] = f'bytes {start}-{stop-1}/{stat_result.st_size}'
                self.headers['Content-Length'] = str(stop - start)
            else:
                self.headers['Content-Range'] = f'bytes {start}-{stat_result.st_size-1}/{stat_result.st_size}'
                self.headers['Content-Length'] = str(
                    stat_result.st_size - start)
        else:
            self.headers['Content-Length'] = str(stat_result.st_size)

        await send(
            {
                "type": "http.response.start",
                "status": self.status_code,
                "headers": self.raw_headers,
            }
        )
        if scope["method"].upper() == "HEAD":
            await send({"type": "http.response.body", "body": b"", "more_body": False})
        else:
            async with aiofiles.open(self.path, mode="rb") as file:
                await file.seek(start)
                more_body = True
                while more_body:
                    old_pos = await file.tell()
                    chunk = await file.read(self.chunk_size)
                    now_pos = await file.tell()
                    if stop is not None and now_pos >= stop:
                        chunk = chunk[:stop-old_pos]
                        more_body = False
                    else:
                        more_body = len(chunk) == self.chunk_size
                    await send({
                        "type": "http.response.body",
                        "body": chunk,
                        "more_body": more_body,
                    })
        if self.background is not None:
            await self.background()


@app.route.get("/download/{any_file_name}")
async def download_file(user: User, path: str, any_file_name: str = ""):
    os = OS(user)
    if (not await os.path.exists(path) or not await os.path.isfile(path)):
        return {
            "code": 1,
            "msg": "文件不存在",
            "data": []
        }
    if (not await os.access(path, os.R_OK)):
        return {
            "code": 1,
            "msg": "权限不足",
            "data": []
        }
    if (user.demo):
        if ((await os.stat(path)).st_size > 1024 * 1024 * 8):
            return {
                "code": 1,
                "msg": "演示模式下禁用本操作",
                "data": []
            }
    return RangeFileResponse(path)

FILE_CHUNK_SIZE = 1024 * 1024*4


@app.route.ws("/upload")
async def new_upload_task(user: User, websocket: WebSocket):
    os = OS(user)
    task_tool = None
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        await websocket.accept()
        manifest = (await websocket.receive_json())
        all_size = 0
        for i in manifest["files"]:
            all_size += i["size"]
        await websocket.send_json({
            "type": "init",
            "data": ""
        })
        if (not await os.access(manifest["dist"], os.W_OK)):
            return {
                "code": 1,
                "msg": "权限不足",
                "data": []
            }
        task_ret_val = True

        async def upload_task(tasktool):
            nonlocal task_tool, task_ret_val
            await (tasktool.write_log("开始上传\n"))
            task_tool = tasktool
            while 1:
                await asyncio.sleep(10000000)
        await execute("tasks", "create_task")(
            ("上传"), f"{len(manifest['files'])} 个文件，总计 {format_size(all_size)}", upload_task, user)
        while True:
            if (task_tool == None):
                await asyncio.sleep(0.1)
            else:
                break
        getbytes = b""
        file_num_id = 0
        now_reveived_size = 0
        all_received_size = 0
        now_file_obj = None
        try:
            if (not await os.path.exists(manifest["dist"]+os.sep+manifest["files"][file_num_id]["path"])):
                await os.makedirs(manifest["dist"]+os.sep +
                                  manifest["files"][file_num_id]["path"])
            add_path = manifest["files"][file_num_id]["path"]
            if (add_path != ""):
                add_path = os.sep+add_path
            now_file_obj = await anyio.open_file(
                manifest["dist"]+add_path+os.sep+manifest["files"][file_num_id]["name"], "wb")
            await task_tool.write_log(
                manifest["files"][file_num_id]["name"]+"\n"
            )  # type: ignore
            while True:
                getbytes = getbytes + await websocket.receive_bytes()
                await websocket.send_json({
                    "type": "next"
                })

                now_reveived_size += len(getbytes)
                if (now_reveived_size >= manifest["files"][file_num_id]["size"]):
                    await now_file_obj.write(getbytes)
                    getbytes = b""
                while len(getbytes) >= FILE_CHUNK_SIZE:
                    await now_file_obj.write(getbytes)
                    getbytes = b""
                if now_reveived_size >= manifest["files"][file_num_id]["size"]:
                    await now_file_obj.aclose()
                    now_file_obj = None
                    file_num_id += 1
                    if file_num_id >= len(manifest["files"]):
                        break
                    all_received_size += now_reveived_size
                    now_reveived_size = 0
                    if (not await os.path.exists(manifest["dist"]+os.sep+manifest["files"][file_num_id]["path"])):
                        await os.makedirs(manifest["dist"]+os.sep +
                                          manifest["files"][file_num_id]["path"])
                    add_path = manifest["files"][file_num_id]["path"]
                    if (add_path != ""):
                        add_path = os.sep+add_path

                    await task_tool.write_log(
                        manifest["files"][file_num_id]["name"]+"\n"
                    )  # type: ignore
                    now_file_obj = await anyio.open_file(
                        manifest["dist"]+add_path+os.sep+manifest["files"][file_num_id]["name"], "wb")
                await task_tool.set_process((now_reveived_size+all_received_size)
                                            * 100/all_size)  # type: ignore
            await task_tool.set_process(100)  # type: ignore
            await task_tool.exit(False)  # type: ignore
            await websocket.close()
        finally:
            if now_file_obj is not None:
                await now_file_obj.aclose()

    except (RuntimeError, websockets.WebSocketDisconnect):
        if (task_tool is not None):
            await task_tool.write_log("用户断开连接\n")  # type: ignore
            await task_tool.exit(True)  # type: ignore


@app.route.get("/attribute")
async def get_file_attribute(user: User, path: str):
    os = OS(user)
    try:
        if (not await os.path.exists(path)):
            return {
                "code": 1,
                "msg": "文件不存在",
                "data": {}
            }
        file_stat = await os.stat(path)
        owner_uid = file_stat.st_uid
        try:
            owner_name = await timeout_second(lambda: pwd.getpwuid(owner_uid).pw_name)
        except KeyError:
            owner_name = "Unknown or Container"
        group_gid = file_stat.st_gid
        try:
            group_name = await timeout_second(lambda: grp.getgrgid(group_gid).gr_name)
        except KeyError:
            group_name = "Unknown or Container"
        file_mode = str(oct(stat.S_IMODE(file_stat.st_mode))
                        ).replace("0o", "").rjust(3, "0")
        if (await os.path.isfile(path)):
            size_bytes = file_stat.st_size
            size_flag = False
            size_str = format_size(size_bytes)
        elif (await os.path.islink(path)):
            size_bytes = 0
            size_flag = False
            size_str = "链接文件(夹)"
        else:
            size_flag = True
            size_bytes = 0
            size_str = "计算中"
        creation_time = await os.path.getctime(path)
        formatted_creation_time = datetime.fromtimestamp(
            creation_time).strftime("%Y-%m-%d %H:%M:%S")
        modification_time = await os.path.getmtime(path)
        formatted_modification_time = datetime.fromtimestamp(
            modification_time).strftime("%Y-%m-%d %H:%M:%S")
    except PermissionError:
        return {
            "code": 1,
            "msg": "权限不足",
            "data": {}
        }

    return {
        "code": 0,
        "msg": "",
        "data": {
            "owner": {
                "name": owner_name,
                "uid": owner_uid
            },
            "group": {
                "name": group_name,
                "gid": group_gid
            },
            "mode": file_mode,
            "size": {
                "str": size_str,
                "bytes": size_bytes,
                "next": size_flag
            },
            "time": {
                "creation": formatted_creation_time,
                "modification": formatted_modification_time
            }
        }
    }


@app.route.post("/attribute/mode")
async def set_file_attribute_mode(user: User, path: str, mode: str):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (not await os.path.exists(path)):
            return {
                "code": 1,
                "msg": "文件不存在"
            }
        await os.chmod(path, int("0o"+mode, 8))
    except PermissionError:
        return {
            "code": 1,
            "msg": "权限不足"
        }
    except Exception as e:
        return {
            "code": 1,
            "msg": repr(e)
        }
    return {
        "code": 0,
        "msg": ""
    }


@app.route.post("/attribute/uid")
async def set_file_attribute_uid(user: User, path: str, uid: int):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (not await os.path.exists(path)):
            return {
                "code": 1,
                "msg": "文件不存在"
            }
        await os.chown(path, int(uid), -1)
    except PermissionError:
        return {
            "code": 1,
            "msg": "权限不足"
        }
    except Exception as e:
        return {
            "code": 1,
            "msg": repr(e)
        }
    return {
        "code": 0,
        "msg": ""
    }


@app.route.post("/attribute/gid")
async def set_file_attribute_gid(user: User, path: str, gid: int):
    os = OS(user)
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    try:
        if (not await os.path.exists(path)):
            return {
                "code": 1,
                "msg": "文件不存在"
            }
        await os.chown(path, -1, int(gid))
    except PermissionError:
        return {
            "code": 1,
            "msg": "权限不足"
        }
    except Exception as e:
        return {
            "code": 1,
            "msg": repr(e)
        }
    return {
        "code": 0,
        "msg": ""
    }


@app.route.get("/search_user_uid", level=2)
async def search_user_uid(user: User, username: str):
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    if (not await timeout_second(lambda: re.match("^[a-zA-Z0-9_]{1,32}$", username))):
        return {
            "code": 1,
            "msg": "用户名不符合规则",
            "data": -1
        }
    try:
        return {
            "code": 0,
            "msg": "",
            "data": await timeout_second(lambda: pwd.getpwnam(username).pw_uid)
        }
    except KeyError:
        return {
            "code": 1,
            "msg": "找不到用户名",
            "data": -1
        }


@app.route.get("/search_group_gid")
async def search_group_gid(user: User, groupname: str):
    if (user.demo):
        return {
            "code": 1,
            "msg": "演示模式下禁用本操作",
            "data": []
        }
    if (not await timeout_second(lambda: re.match("^[a-zA-Z0-9_]{1,32}$", groupname))):
        return {
            "code": 1,
            "msg": "用户名不符合规则",
            "data": -1
        }
    try:
        return {
            "code": 0,
            "msg": "",
            "data": await timeout_second(lambda: grp.getgrnam(groupname).gr_gid)
        }
    except KeyError:
        return {
            "code": 1,
            "msg": "找不到用户名",
            "data": -1
        }


async def get_file_size(os: OS, file_path):
    try:
        ret = (await os.stat(file_path)).st_size
    except PermissionError:
        return 0
    except TimeoutError:
        return 0
    return ret


async def get_dir_size(os: OS, dir_path):
    total_size = 0
    for root, _, files in (await os.walk(dir_path, followlinks=False)):
        file_tasks = [get_file_size(os, os.path.join(root, file))
                      for file in files]
        file_sizes = await asyncio.gather(*file_tasks)
        total_size += sum(file_sizes)

    return total_size


@app.route.ws("/attribute/size")
async def attribute_size(user: User, websocket: WebSocket, path: str):
    os = OS(user)
    search_task = None
    try:
        await websocket.accept()

        async def search_size(path):
            size_sum = await get_dir_size(os, path)
            await websocket.send_json({
                "type": "result",
                "data": {
                    "bytes": size_sum,
                    "str": format_size(size_sum)
                }
            })
            await websocket.close()

        search_task = asyncio.create_task(search_size(path))
        await search_task
        while True:
            await websocket.receive_text()
    except (RuntimeError, websockets.WebSocketDisconnect):
        if (search_task is not None and search_task.done() is False):
            search_task.cancel()
